﻿#include <algorithm>

#include "Stream.h"

_Stream::_Stream(smf::MidiEventList& note_track, double frame_rate, size_t tpq, json& note_sizes) {
	name = _get_name(note_track);
	_tpq = tpq;
	if (note_sizes.contains(name)) {
		_note_frames = (double)note_sizes[name] * frame_rate;
	}
	else {
		_note_frames = 0;
	}
}

std::vector<Message> _Stream::_get_message_list(std::vector<MetaMessage>& meta_track, smf::MidiEventList& note_track, double frame_rate, size_t tpq, size_t note_frames) {
	std::vector<Message> result;
	for (size_t i = 0; i < note_track.size(); i++) {
		Message msg = {
			.start_frame_index = (size_t)round(note_track[i].seconds * frame_rate),
			.end_frame_index = (size_t)round((note_track[i].seconds + note_track[i].getDurationInSeconds()) * frame_rate),
			.velocity = (Byte)note_track[i].getVelocity(),
			.pitch = (Byte)note_track[i].getKeyNumber()
		};
		if (msg.velocity > 127)continue;
		result.push_back(msg);
	}
	std::vector<Message> filtered_result;
	if (note_frames == 0) {
		for (const auto& msg : result) {
			// 添加条件：end_frame_index 不能为初始值 -1ull
			if (msg.velocity > 0 && msg.end_frame_index != (size_t)-1 && msg.end_frame_index > msg.start_frame_index) {
				filtered_result.push_back(msg);
			}
		}
	}
	else {
		for (const auto& msg : result) {
			if (msg.velocity > 0 && msg.end_frame_index != (size_t)-1 && msg.end_frame_index > msg.start_frame_index) {
				Message new_msg = {
					.start_frame_index = msg.start_frame_index,
					.end_frame_index = std::min<size_t>(msg.end_frame_index, msg.start_frame_index + note_frames),
					.velocity = msg.velocity,
					.pitch = msg.pitch
				};
				filtered_result.push_back(new_msg);
			}
		}
	}
	return filtered_result;
}

std::vector<MetaMessage> _Stream::_get_meta_message_list(smf::MidiEventList& meta_track, double frame_rate, size_t tpq) {
	double current_time = 0;
	size_t current_tick = 0;
	size_t current_tempo = 500000;
	std::vector<MetaMessage> result;
	for (size_t i = 0; i < meta_track.size(); i++) {
		smf::MidiEvent msg = meta_track[i];
		current_time += _tick_to_second(msg.tick - current_tick, current_tempo, tpq) * frame_rate;
		current_tick = msg.tick;
		if (msg.isTempo()) {
			current_tempo = msg.getTempoMicroseconds();
			MetaMessage meta = {
				.frame_index = (size_t)round(current_time),
				.tick_index = current_tick,
				.tempo = current_tempo
			};
			result.push_back(meta);
		}
	}
	return result;
}

size_t _Stream::_get_track_frames_num(std::vector<MetaMessage>& meta_track, smf::MidiEventList& note_track, double frame_rate, size_t tpq) {
	double current_time = 0;
	for (size_t i = 0; i < note_track.size(); i++) {
		double value = note_track[i].seconds + note_track[i].getDurationInSeconds();
		if (current_time < value)current_time = value;
	}
	return (size_t)round(current_time * frame_rate);
}

bool _Stream::_compare_tick(CombinedMessage& a, CombinedMessage& b) {
	return a.current_tick < b.current_tick;
}

std::string _Stream::_get_name(smf::MidiEventList& note_track) {
	for (size_t i = 0; i < note_track.size(); i++) {
		const smf::MidiEvent& event = note_track[i];

		// 检查是否为音轨名称事件 (FF 03)
		if (event.isTrackName() && event.size() >= 3) {
			// 事件结构: [FF][03][VLQ长度][名称数据...]
			size_t pos = 2; // 跳过 FF 和 03
			uint32_t length = 0;

			// 正确解析 VLQ 长度
			do {
				if (pos >= event.size()) return ""; // 防止越界
				length = (length << 7) | (event[pos] & 0x7F);
			} while ((event[pos++] & 0x80) != 0 && pos < 4); // VLQ 最多 4 字节

			// 提取名称数据
			if (pos + length > event.size()) return ""; // 长度校验
			const uint8_t* name_data = event.data() + pos;
			return std::string(name_data, name_data + length);
		}
	}
	return ""; // 未找到音轨名称
}

double _Stream::_tick_to_second(size_t tick, size_t tempo, size_t tpq) {
	return 1e-6 * tempo / tpq * tick;
}

NoteStream::NoteStream() {
	pitch_min = -1;
	pitch_range = -1;
	notes_num = -1;
}

NoteStream::NoteStream(smf::MidiEventList& meta_track, smf::MidiEventList& note_track, double frame_rate, size_t tpq, json& note_sizes) :_Stream(note_track, frame_rate, tpq, note_sizes) {
	notes = _get_stream_data(meta_track, note_track, frame_rate, tpq, _note_frames);
	pitch_min = _get_pitch_min(notes);
	pitch_range = _get_pitch_range(notes, pitch_min);
	notes_num = notes.rows;
}

void NoteStream::align(size_t begin_frames) {
	for (size_t i = 0; i < notes.rows; i++) {
		notes.at(i, 0) += (double)begin_frames;
		notes.at(i, 1) += (double)begin_frames;
	}
}

NoteStream NoteStream::copy() {
	NoteStream new_stream;
	new_stream.name = name;
	new_stream._note_frames = _note_frames;
	new_stream._tpq = _tpq;
	new_stream.notes = notes.clone();
	new_stream.pitch_min = pitch_min;
	new_stream.pitch_range = pitch_range;
	new_stream.notes_num = notes_num;
	return new_stream;
}

size_t NoteStream::_get_pitch_min(Mat& data) {
	size_t pitch_min = 127;
	for (size_t i = 0; i < data.rows; i++) {
		size_t pitch = (size_t)data.at(i, 3);
		if (pitch < pitch_min)pitch_min = pitch;
	}
	return pitch_min;
}

size_t NoteStream::_get_pitch_range(Mat& data, size_t pitch_min) {
	size_t pitch_max = 0;
	for (size_t i = 0; i < data.rows; i++) {
		size_t pitch = (size_t)data.at(i, 3);
		if (pitch > pitch_max)pitch_max = pitch;
	}
	return pitch_max - pitch_min + 1;
}

Mat NoteStream::_get_stream_data(smf::MidiEventList& meta_track, smf::MidiEventList& note_track, double frame_rate, size_t tpq, size_t note_frames) {
	std::vector<MetaMessage> meta_data = _get_meta_message_list(meta_track, frame_rate, tpq);
	std::vector<Message> message_data = _get_message_list(meta_data, note_track, frame_rate, tpq, note_frames);
	Mat notes_data = Mat(message_data.size(), 4, 0);
	size_t message_count = 0;
	for (size_t i = 0; i < message_data.size(); i++) {
		Message message = message_data[i];
		if (message.start_frame_index == message.end_frame_index || message.velocity == 0) continue;
		notes_data.at(i, 0) = (double)message.start_frame_index;
		notes_data.at(i, 1) = (double)message.end_frame_index;
		notes_data.at(i, 2) = (double)message.velocity;
		notes_data.at(i, 3) = (double)message.pitch;
		message_count++;
	}
	meta_data.clear();
	message_data.clear();
	Mat notes_data2 = notes_data(0, message_count, 0, 4);
	return notes_data2;
}

VelocityStream::VelocityStream() {
	_blur_sigma = -1;
	length = -1;
}

VelocityStream::VelocityStream(smf::MidiEventList& meta_track, smf::MidiEventList& note_track, double frame_rate, double blur_sigma, size_t tpq, json& note_sizes) :_Stream(note_track, frame_rate, tpq, note_sizes) {
	velocity = _get_stream_data(meta_track, note_track, frame_rate, tpq, _note_frames);
	_blur_sigma = blur_sigma;
	length = velocity.rows;
}

void VelocityStream::align(size_t total_frame_num, size_t begin_frames, size_t end_frames) {
	if (velocity.rows < total_frame_num) {
		Mat temp_mat = velocity.clone();
		size_t zeros_length = total_frame_num - velocity.rows;
		Mat zeros_velocity = Mat(zeros_length, 1, 0);
		Mat new_mat = vconcat(temp_mat, zeros_velocity);
		velocity = new_mat.clone();
		new_mat.release();
		length += zeros_length;
	}
	Mat temp_mat = velocity.clone();
	Mat begin_zeros_velocity = Mat(begin_frames, 1, 0);
	Mat end_zeros_velocity = Mat(end_frames, 1, 0);
	Mat new_mat = vconcat(begin_zeros_velocity, temp_mat);
	velocity = vconcat(new_mat, end_zeros_velocity);
	new_mat.release();
	temp_mat.release();
	begin_zeros_velocity.release();
	end_zeros_velocity.release();
	length += begin_frames + end_frames;
}

void VelocityStream::clear() {
	velocity.release();
}

VelocityStream VelocityStream::copy() {
	VelocityStream new_stream;
	new_stream.name = name;
	new_stream._note_frames = _note_frames;
	new_stream._tpq = _tpq;
	new_stream.length = length;
	new_stream.velocity = velocity.clone();
	new_stream._blur_sigma = _blur_sigma;
	return new_stream;
}

void VelocityStream::gaussian_blur() {
	velocity.blur(_blur_sigma);
}

double VelocityStream::get(size_t frame_index) {
	return velocity.at(frame_index, 0);
}

Mat VelocityStream::_get_stream_data(smf::MidiEventList& meta_track, smf::MidiEventList& note_track, double frame_rate, size_t tpq, size_t note_frames) {
	std::vector<MetaMessage> meta_data = _get_meta_message_list(meta_track, frame_rate, tpq);
	size_t _length = _get_track_frames_num(meta_data, note_track, frame_rate, tpq);
	Mat velocity_data = Mat(_length, 1, 0);
	std::vector<Message> message_data = _get_message_list(meta_data, note_track, frame_rate, tpq, note_frames);
	for (size_t i = 0; i < message_data.size(); i++) {
		Message message = message_data[i];
		for (size_t j = message.start_frame_index; j < message.end_frame_index; j++) {
			double value = velocity_data.at(j, 0);
			if (value < message.velocity)velocity_data.at(j, 0) = (double)message.velocity;
		}
	}
	meta_data.clear();
	message_data.clear();
	return velocity_data;
}
